import json

from app.schema.base import RespBaseSchema, RespDataSchema
from app.service.base import BaseService
from app.dao.message import MessageDao
from app.dao.ai import AiDao
from app.dao.chat import ChatDao
from app.dao.user import UserDao
from app.model.message import Message
from app.schema.message import *
from app.schema.base import *
from app.chain.chat_chain import ChatDocsChain


class MessageService(BaseService):
    def __init__(self, auth_data: dict = {}):
        user_id = auth_data.id
        self.Model = Message
        self.dao = MessageDao(user_id)
        self.dao.Model = Message

        super().__init__(user_id, auth_data)

    async def send_one_message(self, message: MessageInfoSchema): #保存消息到数据库,做出请求流式响应,在结束后把响应结果保存到数据库
        res = self.create(message)
        chat_info = ChatDao().read(message.chat_id)
        if chat_info:
            ai_info = AiDao().read(chat_info.ai_id)
            history = await self.get_history_func(message.chat_id)
            chain=ChatDocsChain(ai_info.vector_name)
            ans = ''
            yield {"event":"response"}
            async for token in chain.chat(query=f"{message.content}\n", history=history, ai_prompt=ai_info.remark):
                ans += token
                yield {"event":"msg","token":token}
            ai_message = MessageInfoSchema(user_id = 2, chat_id = message.chat_id, content = ans, user_type = 1, message_type=0, parent_id=res.id)
            self.create(ai_message)
        else:
            raise RespBaseSchema(code=400, message='发送失败')

    def get_message_list(self, args:ListArgsSchema):
        args.keys = [
            ListKeySchema(key='id'),
            ListKeySchema(key='message_type'),
            ListKeySchema(key='created_time',rename='created_time'),
            ListKeySchema(key='content'),
            ListKeySchema(key='user_type'),
        ]
        args.orders =[
            ListOrderSchema(key='id',condition='desc')]
        
        res = self.list(args)
        for item in res.list:
            if item["message_type"] == 1:
                chat_info = ChatDao().read(args.chat_id)
                if chat_info:
                    ai_info = AiDao().read(chat_info.ai_id)
                    user_info = UserDao().read(chat_info.group_id)
                    item["content"] = json.dumps({
                        "name":user_info.name,
                        "avatar":user_info.avatar,
                        "info":ai_info.info
                    })
            elif item["message_type"] == 2:
                chat_info = ChatDao().read(args.chat_id)
                if chat_info:
                    ai_info = AiDao().read(chat_info.ai_id)
                    item["content"] = ai_info.greet

        return res
    
    async def get_history_func(self, chat_id: int):
        args = ListArgsSchema(chat_id=chat_id)
        args.size = 11
        args.keys = [
            ListKeySchema(key='id'),
            ListKeySchema(key='content'),
            ListKeySchema(key='user_type'),
            ListKeySchema(key='parent_id')
        ]
        args.orders =[ListOrderSchema(key='created_time',condition='desc')]
        
        history = self.list(args).list

        ans = []

        for mes in history[0:3]:
            if mes['parent_id'] is not None:
                parent_mes = self.read(mes['parent_id'])
                if parent_mes:
                    temp = {
                        "id":parent_mes.id,
                        "content":parent_mes.content,
                        "user_type":parent_mes.user_type
                    }
                    ans.append(mes)
                    ans.append(temp)
        ans.reverse()
        return ans
    
    async def stream(self, message: MessageInfoSchema):
        chat_info = ChatDao().read(message.chat_id)
        if chat_info:
            ai_info = AiDao().read(chat_info.ai_id)
            history = []
            chain=ChatDocsChain(ai_info.vector_name)
            async for token in chain.chat_content(query=f"{message.content}\n", history=history, ai_prompt=ai_info.remark):
                yield {"event":"msg","token":token}

        else:
            raise RespBaseSchema(code=400, message='发送失败')
        
    async def nostream(self, message: MessageInfoSchema):
        chat_info = ChatDao().read(message.chat_id)
        if chat_info:
            ai_info = AiDao().read(chat_info.ai_id)
            history = []
            chain=ChatDocsChain(ai_info.vector_name)
            ans = ''
            async for token in chain.chat_content(query=f"{message.content}\n", history=history, ai_prompt=ai_info.remark):
                ans += token
            return ans
        else:
            raise RespBaseSchema(code=400, message='发送失败')